#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import numpy as np
from typing import Dict, List, Tuple
import logging
from datetime import datetime
import json
import os
from scapy.all import sniff, IP, TCP, UDP, TLS

from .vpn_knowledge import VPNProtocol, ProtocolFeature, VPN_KNOWLEDGE_BASE
from .rag_detector import RAGVPNDetector

class TrainingDataGenerator:
    """VPN协议训练数据生成器"""
    
    def __init__(self, output_dir: str = "training_data"):
        self.logger = logging.getLogger(__name__)
        self.output_dir = output_dir
        self.detector = RAGVPNDetector()
        
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
            
    def capture_traffic(self, interface: str, duration: int = 300):
        """捕获网络流量"""
        self.logger.info(f"开始在接口 {interface} 上捕获流量，持续 {duration} 秒")
        
        packets = []
        def packet_callback(packet):
            if packet.haslayer(IP):
                packets.append(packet)
                
        sniff(iface=interface, prn=packet_callback, timeout=duration)
        return packets
        
    def generate_training_data(self, packets: List[Packet], protocol: VPNProtocol):
        """生成训练数据"""
        self.logger.info(f"开始为协议 {protocol.value} 生成训练数据")
        
        # 1. 提取特征
        features = []
        for packet in packets:
            result = self.detector.detect(packet)
            if result.protocol == protocol:
                features.append({
                    'packet_size': len(packet),
                    'protocol': protocol.value,
                    'timestamp': datetime.now().isoformat(),
                    'flow_features': result.features,
                    'evidence': result.evidence
                })
                
        # 2. 保存训练数据
        output_file = os.path.join(self.output_dir, f"{protocol.value.lower()}_features.json")
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(features, f, indent=2, ensure_ascii=False)
            
        self.logger.info(f"已保存 {len(features)} 条训练数据到 {output_file}")
        
    def update_protocol_features(self, protocol: VPNProtocol):
        """更新协议特征"""
        self.logger.info(f"开始更新协议 {protocol.value} 的特征")
        
        # 1. 加载训练数据
        data_file = os.path.join(self.output_dir, f"{protocol.value.lower()}_features.json")
        if not os.path.exists(data_file):
            self.logger.warning(f"未找到协议 {protocol.value} 的训练数据")
            return
            
        with open(data_file, 'r', encoding='utf-8') as f:
            features = json.load(f)
            
        # 2. 分析特征
        packet_sizes = [f['packet_size'] for f in features]
        flow_features = [f['flow_features'] for f in features]
        
        # 3. 更新知识库
        current_feature = VPN_KNOWLEDGE_BASE[protocol]
        updated_feature = ProtocolFeature(
            name=protocol,
            ports=current_feature.ports,  # 保持端口不变
            packet_patterns=current_feature.packet_patterns,  # 保持包模式不变
            tls_fingerprints=current_feature.tls_fingerprints,  # 保持TLS指纹不变
            flow_characteristics={
                'min_packet_size': np.percentile(packet_sizes, 5),
                'max_packet_size': np.percentile(packet_sizes, 95),
                'protocol_ratio': self._calculate_protocol_ratio(flow_features)
            },
            detection_difficulty=current_feature.detection_difficulty,
            evasion_techniques=current_feature.evasion_techniques
        )
        
        # 4. 保存更新后的特征
        VPN_KNOWLEDGE_BASE[protocol] = updated_feature
        self.logger.info(f"已更新协议 {protocol.value} 的特征")
        
    def _calculate_protocol_ratio(self, flow_features: List[Dict]) -> Dict[str, float]:
        """计算协议比例"""
        protocol_counts = defaultdict(int)
        total_flows = len(flow_features)
        
        for features in flow_features:
            for proto in ['TCP', 'UDP', 'ESP']:
                if features.get(f'{proto.lower()}_ratio', 0) > 0:
                    protocol_counts[proto] += 1
                    
        return {
            proto: count / total_flows
            for proto, count in protocol_counts.items()
        }
        
    def generate_all_protocols_data(self, interface: str, duration: int = 300):
        """为所有协议生成训练数据"""
        for protocol in VPNProtocol:
            self.logger.info(f"开始为协议 {protocol.value} 生成训练数据")
            packets = self.capture_traffic(interface, duration)
            self.generate_training_data(packets, protocol)
            self.update_protocol_features(protocol) 